package com.peng.leetcode.tree;

import java.util.HashMap;
import java.util.Map;

/**
 * PrefixTree
 * 剑指 Offer II 062. 实现前缀树
 * https://leetcode.cn/problems/QC3q1f/
 * <p>
 * Created on 2022/11/19
 *
 * @author lupeng
 */
public class PrefixTree {

    class Trie {

        static class LevelNode {
            Map<Character, LevelNode> next = new HashMap<>();
            boolean end = false;
        }

        LevelNode root = new LevelNode();

        /**
         * Initialize your data structure here.
         */
        public Trie() {
        }

        /**
         * Inserts a word into the trie.
         */
        public void insert(String word) {
            LevelNode curr = root;
            int length = word.length();
            for (int i = 0; i < length; i++) {
                Character c = word.charAt(i);
                if (!curr.next.containsKey(c)) {
                    curr.next.put(c, new LevelNode());
                }
                curr = curr.next.get(c);
            }
            curr.end = true;
        }

        /**
         * Returns if the word is in the trie.
         */
        public boolean search(String word) {
            LevelNode curr = root;
            int length = word.length();
            for (int i = 0; i < length; i++) {
                Character c = word.charAt(i);
                if (!curr.next.containsKey(c)) {
                    return false;
                }
                curr = curr.next.get(c);
            }
            return curr.end;
        }

        /**
         * Returns if there is any word in the trie that starts with the given prefix.
         */
        public boolean startsWith(String prefix) {
            LevelNode curr = root;
            int length = prefix.length();
            for (int i = 0; i < length; i++) {
                Character c = prefix.charAt(i);
                if (!curr.next.containsKey(c)) {
                    return false;
                }
                curr = curr.next.get(c);
            }
            return true;
        }
    }
}
